const { NlpManager } = require('node-nlp');
const { pipeline } = require('@xenova/transformers');
const fs = require('fs');

// 配置镜像源（在代码最前面添加）
process.env.TRANSFORMERS_CACHE = './models'; // 修改缓存路径
process.env.HF_ENDPOINT = 'https://hf-mirror.com'; // 使用国内镜像

// 初始化 NLP 管理器（用于规则匹配）
const manager = new NlpManager({ languages: ['zh'] });

// 从JSON文件加载知识库
const knowledgeBase = JSON.parse(fs.readFileSync( __dirname+  '/knowledge.json'));

// 在训练时动态添加问答
knowledgeBase.forEach(item => {
  manager.addDocument('zh', item.question, item.intent);
  manager.addAnswer('zh', item.intent, item.answer);
});

// 2. 初始化生成式模型（用于未见过的问题）
let generativeModel;
async function loadModel() {
    console.log('正在加载AI模型...（首次运行需要下载模型）');
    generativeModel = await pipeline('text-generation', 'Xenova/gpt2-small-chinese');
    console.log('模型加载完成！');
}

// 3. 混合回答函数
async function getAnswer(question) {
    // 先尝试用预设规则回答
    const nlpResult = await manager.process('zh', question);
    
    if (nlpResult.answer && nlpResult.score > 0.7) {
        return { 
            answer: nlpResult.answer,
            type: '规则匹配',
            confidence: nlpResult.score.toFixed(2)
        };
    } else {
        // 没有匹配时调用生成模型
        const generated = await generativeModel(question, {
            max_length: 100,
            num_return_sequences: 1
        });
        return {
            answer: generated[0].generated_text.replace(question, '').trim(),
            type: 'AI生成',
            confidence: 'N/A'
        };
    }
}

// 4. 命令行交互
const readline = require('readline').createInterface({
    input: process.stdin,
    output: process.stdout
});

async function main() {
    await manager.train();
    await loadModel();
    
    console.log('聊天机器人已启动（输入"退出"结束）\n');
    
    async function ask() {
        readline.question('> ', async (question) => {
            if (question.toLowerCase() === '退出') {
                readline.close();
                return;
            }
            
            const startTime = Date.now();
            const response = await getAnswer(question);
            const timeTaken = ((Date.now() - startTime) / 1000).toFixed(2);
            
            console.log('\n回答:', response.answer);
            console.log(`[类型: ${response.type} | 置信度: ${response.confidence} | 耗时: ${timeTaken}s]\n`);
            ask();
        });
    }
    
    ask();
}

main().catch(console.error);